from langchain.agents import Tool
from langchain_community.tools.file_management.read import ReadFileTool
from langchain_community.tools.file_management.write import WriteFileTool
from langchain_community.utilities import SerpAPIWrapper
from langchain_openai import OpenAIEmbeddings
from autogpt.agent import AutoGPT
from langchain_openai import ChatOpenAI
from langchain_community.embeddings import QianfanEmbeddingsEndpoint
from langchain_community.llms import QianfanLLMEndpoint
import os
from langchain_core.tools import tool

os.environ["SERPAPI_API_KEY"] = "5f6773509c10d3232df4ccc04f2636a647065900a4a1f3c5dbd05d485820c991"
os.environ["QIANFAN_AK"] = "7ZITkqmufR2g6MSLQUqsYAXB"
os.environ["QIANFAN_SK"] = "n2YLOoakGri5LfsK8Fxsax86tdeLEfpm"
@tool
def external_user_list(query:str) -> list[str]:
    """返回客户列表"""
    l = ['培生', '文龙', '小明']
    print(f'external_user_list {l}')
    return l

@tool
def send_message(user_id:str, message:str) -> bool:
    """给客户发消息"""
    print(f'send_message {message} to user:{user_id}')
    return True

search = SerpAPIWrapper()
tools = [
    Tool(
        name="search",
        func=search.run,
        description="useful for when you need to answer questions about current events. You should ask targeted questions",
    ),
    external_user_list,
    send_message,
    WriteFileTool(),
    ReadFileTool(),
]
api_key = "34a3cc1f-3608-4c2e-8cbe-6f2076c30126"
model = "ep-20240606130022-nrz5h"
# Define your embedding model
embeddings = QianfanEmbeddingsEndpoint(
    chunk_size=400
)
# Initialize the vectorstore as empty
import faiss
from langchain_community.vectorstores import Chroma
vectorstore = Chroma(embedding_function=embeddings)

model = ChatOpenAI(openai_api_key=api_key,
                     openai_api_base="https://ark.cn-beijing.volces.com/api/v3",
                     model=model,
                     temperature=0.8)
agent = AutoGPT.from_llm_and_tools(
    ai_name="Tom",
    ai_role="Assistant",
    tools=tools,
    llm=model,
    memory=vectorstore.as_retriever(),
)
# Set verbose to be true
agent.chain.verbose = True

agent.run(["做一个群发活动，步骤如下：1，找到客户列表；2，给每个客户生成节日祝福文案；3，给每个客户发送祝福消息"])
